package ast

import (
	"fmt"
	"strings"

	"gitee.com/yanxingshuyuan/ming/token"
)

type Node interface {
	TokenLiteral() string
	String() string
}

type Expr interface {
	Node
	exprNode()
}

type Stmt interface {
	Node
	stmtNode()
}

// 表示一段程序
type Program struct {
	Node
	Stmts []Stmt
}

func (p *Program) TokenLiteral() string {
	if len(p.Stmts) > 0 {
		return p.Stmts[0].TokenLiteral()
	}
	return ""
}

func (p *Program) String() string {
	var out strings.Builder
	for _, s := range p.Stmts {
		out.WriteString(s.String())
	}
	return out.String()
}

type VarStmt struct {
	Token token.Token
	Name  *Ident
	Value Expr
}

func (vs *VarStmt) stmtNode()            {}
func (vs *VarStmt) TokenLiteral() string { return string(vs.Token.Literal) }
func (vs *VarStmt) String() string {
	var out strings.Builder
	out.WriteString("var ")
	out.WriteString(vs.Name.String())
	if vs.Value != nil {
		out.WriteString(" = ")
		out.WriteString(vs.Value.String())
	}
	out.WriteString("\n")
	return out.String()
}

type ExprStmt struct {
	Token token.Token
	Expr  Expr
}

func (es *ExprStmt) stmtNode()            {}
func (es *ExprStmt) TokenLiteral() string { return string(es.Token.Literal) }
func (es *ExprStmt) String() string {
	if es.Expr != nil {
		return es.Expr.String()
	}
	return ""
}

type IntLiteral struct {
	Token token.Token
	Value int
}

func (il *IntLiteral) exprNode()            {}
func (il *IntLiteral) TokenLiteral() string { return string(il.Token.Literal) }
func (il *IntLiteral) String() string       { return fmt.Sprintf("%d", il.Value) }

type BoolLiteral struct {
	Token token.Token
	Value bool
}

func (bl *BoolLiteral) exprNode()            {}
func (bl *BoolLiteral) TokenLiteral() string { return string(bl.Token.Literal) }
func (bl *BoolLiteral) String() string       { return fmt.Sprintf("%t", bl.Value) }

type Ident struct {
	Token token.Token
	Value string
}

func (id *Ident) exprNode()            {}
func (id *Ident) TokenLiteral() string { return string(id.Token.Literal) }
func (id *Ident) String() string       { return id.Value }

type PrefixExpr struct {
	Token    token.Token
	Operator string
	Right    Expr
}

func (pe *PrefixExpr) exprNode()            {}
func (pe *PrefixExpr) TokenLiteral() string { return string(pe.Token.Literal) }
func (pe *PrefixExpr) String() string {
	var out strings.Builder

	out.WriteString("(")
	out.WriteString(pe.Operator)
	out.WriteString(pe.Right.String())
	out.WriteString(")")
	return out.String()
}

type InfixExpr struct {
	Token    token.Token
	Operator string
	Left     Expr
	Right    Expr
}

func (ie *InfixExpr) exprNode()            {}
func (ie *InfixExpr) TokenLiteral() string { return string(ie.Token.Literal) }
func (ie *InfixExpr) String() string {
	var out strings.Builder

	out.WriteString("(")
	out.WriteString(ie.Left.String())
	out.WriteString(" " + ie.Operator + " ")
	out.WriteString(ie.Right.String())
	out.WriteString(")")
	return out.String()
}

type IfExpr struct {
	Token     token.Token
	Condition Expr
	Then      *BlockStmt
	Else      *BlockStmt
}

func (ie *IfExpr) exprNode()            {}
func (ie *IfExpr) TokenLiteral() string { return string(ie.Token.Literal) }
func (ie *IfExpr) String() string {
	var out strings.Builder

	out.WriteString("if")
	out.WriteString(ie.Condition.String())
	out.WriteString(" ")
	out.WriteString(ie.Then.String())

	if ie.Else != nil {
		out.WriteString(" else ")
		out.WriteString(ie.Else.String())
	}

	return out.String()
}

type BlockStmt struct {
	Token token.Token
	Stmts []Stmt
}

func (bs *BlockStmt) stmtNode()            {}
func (bs *BlockStmt) TokenLiteral() string { return string(bs.Token.Literal) }
func (bs *BlockStmt) String() string {
	var out strings.Builder
	out.WriteString("{\n")
	if bs != nil {
		for _, s := range bs.Stmts {
			out.WriteString(s.String())
		}
	}
	out.WriteString("\n}\n")
	return out.String()
}

type ForExpr struct {
	Token     token.Token
	Initial   Stmt
	Condition Expr
	Step      Stmt
	Body      *BlockStmt
}

func (fe *ForExpr) exprNode()            {}
func (fe *ForExpr) TokenLiteral() string { return string(fe.Token.Literal) }
func (fe *ForExpr) String() string {
	var out strings.Builder

	out.WriteString("for ")
	if fe.Initial != nil {
		out.WriteString(fe.Initial.String())
	}
	out.WriteString("; ")
	out.WriteString(fe.Condition.String())
	out.WriteString("; ")
	if fe.Initial != nil {
		out.WriteString(fe.Step.String())
	}
	out.WriteString(fe.Body.String())

	return out.String()
}

type FnDeclStmt struct {
	Token      token.Token
	Name       *Ident
	Params     []*Param
	ReturnType Type
	Body       *BlockStmt
}

func (fds *FnDeclStmt) stmtNode()            {}
func (fds *FnDeclStmt) TokenLiteral() string { return string(fds.Token.Literal) }
func (fds *FnDeclStmt) String() string {
	var out strings.Builder
	out.WriteString("func ")
	out.WriteString(fds.Name.String())
	out.WriteString("(")
	for _, p := range fds.Params {
		out.WriteString(p.String())
		out.WriteString(", ")
	}
	out.WriteString(")")

	out.WriteString(string(fds.ReturnType))
	out.WriteString(fds.Body.String())

	return out.String()
}

type Param struct {
	Name *Ident
	Type Type
}

func (p *Param) String() string {
	var out strings.Builder

	out.WriteString(p.Name.String())
	out.WriteString(" ")
	out.WriteString(string(p.Type))

	return out.String()
}

type Type string

const (
	Void Type = "void"
	Int       = "int"
	Bool      = "bool"
)

type CallExpr struct {
	Token token.Token
	Fn    Expr
	Args  []Expr
}

func (ce *CallExpr) exprNode()            {}
func (ce *CallExpr) TokenLiteral() string { return string(ce.Token.Literal) }
func (ce *CallExpr) String() string {
	var out strings.Builder

	args := []string{}
	for _, a := range ce.Args {
		args = append(args, a.String())
	}

	out.WriteString(ce.Fn.String())
	out.WriteString("(")
	out.WriteString(strings.Join(args, ", "))
	out.WriteString(")")

	return out.String()
}
